package se.webbzon.oschi01.aes;


public class RijndaelMatrix {
	
	// The byte array containing the Rijndael matrix.
	private final byte[] state;
	
	// The offset for the state matrix in the byte array
	private final int offset;
	
	/*============================================================
	Constructors
	============================================================*/
	
	/** Creates a new empty Rijndael matrix. **/
	private RijndaelMatrix() {
		this(new byte[16],0);
	}
	
	/** Creates a new Rijndael matrix which occupies 16 bytes in
	 * the provided byte array starting at the given offset. **/
	private RijndaelMatrix(byte[] data, int offset) {
		state = data;
		this.offset = offset;
	}
	
	/*============================================================
	Public Methods
	============================================================*/
	
	/** Creates an ordered array of Rijndael matrices which contains
	 * the information of the provided byte array. The provided data
	 * array must have a length which is a multiple of 16. Any changes 
	 * in the provided data array will affect the returned matrices. **/
	public static RijndaelMatrix[] fromBytes(byte[] data, int offset, int length) {
		if (length % 16 != 0)
			throw new IllegalArgumentException("Data array must be a multiple of 16. Length = " + length);
		final RijndaelMatrix[] matrices = new RijndaelMatrix[length/16];
		final int limit = length / 16;
		for (int i = 0; i < limit; i++) {
			final int o = offset + i * 16;
			matrices[i] = new RijndaelMatrix(data,o);
		}
		return matrices;
	}
	
	/** Creates a new Rijndael matrix which can be used as a key when 
	 * encrypting messages. Any changes in the provided data array will
	 * affect the key. **/
	public static RijndaelMatrix createKey(byte[] data, int offset) {
		if (data.length - offset < 16)
			throw new IllegalArgumentException("The remaining data in the array is smaller than 16 bytes. Remaining = " + (data.length - offset));
		return new RijndaelMatrix(data,offset);
	}
	
	/** Performs the subBytes step on this Rijndael state by using the provided
	 * s-Box. **/
	public final void subBytes(RijndaelSBox sBox) {
		for (int i = 0; i < 16; i++)
			set(i,sBox.swap(get(i)));
	}
	
	/** Performs the inverse of the subBytes step on this Rijndael state by 
	 * using the provided s-Box. **/
	public final void invSubBytes(RijndaelSBox sBox) {
		for (int i = 0; i < 16; i++)
			set(i,sBox.swapBack(get(i)));
	}
	
	/** Performs the shiftRows step on this Rijndael state. **/
	public final void shiftRows() {
		// Second row
		byte a = get(1);
		set(1, get(5));
		set(5, get(9));
		set(9, get(13));
		set(13, a);
		// Third row
		a = get(2);
		set(2, get(10));
		set(10, a);
		a = get(6);
		set(6, get(14));
		set(14, a);
		// Fourth row;
		a = get(15);
		set(15, get(11));
		set(11, get(7));
		set(7, get(3));
		set(3, a);
	}
	
	/** Performs the inverse of the shiftRows step on this Rijndael state. **/
	public final void invShiftRows() {
		// Second row
		byte a = get(13);
		set(13, get(9));
		set(9, get(5));
		set(5, get(1));
		set(1, a);
		// Third row
		a = get(2);
		set(2, get(10));
		set(10, a);
		a = get(6);
		set(6, get(14));
		set(14, a);
		// Fourth row;
		a = get(3);
		set(3, get(7));
		set(7, get(11));
		set(11, get(15));
		set(15, a);
	}
	
	/** Performs the mixColumns step on this Rijndael state. **/
	public final void mixColumns() {
		final byte c2 = (byte) 0x02, c3 = (byte) 0x03;
		byte[] tmpState = new byte[16];
		for (int i = 0; i < 16; i+=4) {
			tmpState[i] 	= (byte) (	gMul(c2, 	get(i)) ^ 
										gMul(c3, 	get(i+1)) ^ 
													get(i+2) ^ 
													get(i+3));
			tmpState[i+1] 	= (byte) (				get(i) ^ 
										gMul(c2, 	get(i+1)) ^ 
										gMul(c3, 	get(i+2)) ^ 
													get(i+3));
			tmpState[i+2]	= (byte) (				get(i) ^ 
													get(i+1) ^ 
										gMul(c2, 	get(i+2)) ^ 
										gMul(c3, 	get(i+3)));
			tmpState[i+3]	= (byte) (	gMul(c3, 	get(i)) ^ 
													get(i+1) ^ 
													get(i+2) ^ 
										gMul(c2, 	get(i+3)));
		}
		
		// Copy the temporary state to the actual state
		for (int i = 0; i < 16; i++)
			set(i, tmpState[i]);
	}
	
	/** Performs the inverse of the mixColumns step on this Rijndael state. **/
	public final void invMixColumns() {
		final byte 	c9 = (byte) 0x09, c11 = (byte) 0x0b,
					c13 = (byte) 0x0d, c14 = (byte) 0x0e;
		byte[] tmpState = new byte[16];
		for (int i = 0; i < 16; i+=4) {
			tmpState[i] 	= (byte) (	gMul(c14, 	get(i)) ^ 
										gMul(c11, 	get(i+1)) ^ 
										gMul(c13,	get(i+2)) ^ 
										gMul(c9,	get(i+3)));
			tmpState[i+1] 	= (byte) (	gMul(c9, 	get(i)) ^ 
										gMul(c14, 	get(i+1)) ^ 
										gMul(c11,	get(i+2)) ^ 
										gMul(c13,	get(i+3)));
			tmpState[i+2]	= (byte) (	gMul(c13, 	get(i)) ^ 
										gMul(c9, 	get(i+1)) ^ 
										gMul(c14,	get(i+2)) ^ 
										gMul(c11,	get(i+3)));
			tmpState[i+3]	= (byte) (	gMul(c11, 	get(i)) ^ 
										gMul(c13, 	get(i+1)) ^ 
										gMul(c9,	get(i+2)) ^ 
										gMul(c14,	get(i+3)));
		}
		
		// Copy the temporary state to the actual state
		for (int i = 0; i < 16; i++)
			set(i, tmpState[i]);
	}
	
	/** Performs the addRoundKey step on this Rijndael state by using the
	 * given sub-key. **/
	public final void addRoundKey(RijndaelMatrix subKey) {
		for (int i = 0; i < 16; i++)
			set(i, (byte) (get(i) ^ subKey.get(i)));
	}
	
	/** Creates a new sub-key by using this Rijndael key and the rcon 
	 * input i. **/
	public final RijndaelMatrix createSubKey(RijndaelSBox sBox, byte i) {
		RijndaelMatrix subKey = new RijndaelMatrix();
		// Generate first column
		byte x1 = sBox.swap(get(13));
		byte x2 = sBox.swap(get(14));
		byte x3 = sBox.swap(get(15));
		byte x4 = sBox.swap(get(12));
		byte r = rcon(i);
		subKey.set(0, (byte) (get(0) ^ x1 ^ r));
		subKey.set(1, (byte) (get(1) ^ x2));
		subKey.set(2, (byte) (get(2) ^ x3));
		subKey.set(3, (byte) (get(3) ^ x4));
		// Generate second column
		subKey.set(4, (byte) (subKey.get(0) ^ get(4)));
		subKey.set(5, (byte) (subKey.get(1) ^ get(5)));
		subKey.set(6, (byte) (subKey.get(2) ^ get(6)));
		subKey.set(7, (byte) (subKey.get(3) ^ get(7)));
		// Generate third column
		subKey.set(8, (byte) (subKey.get(4) ^ get(8)));
		subKey.set(9, (byte) (subKey.get(5) ^ get(9)));
		subKey.set(10, (byte) (subKey.get(6) ^ get(10)));
		subKey.set(11, (byte) (subKey.get(7) ^ get(11)));
		// Generate fourth column
		subKey.set(12, (byte) (subKey.get(8) ^ get(12)));
		subKey.set(13, (byte) (subKey.get(9) ^ get(13)));
		subKey.set(14, (byte) (subKey.get(10) ^ get(14)));
		subKey.set(15, (byte) (subKey.get(11) ^ get(15)));
		return subKey;
	}
	
	@Override public String toString() {
		StringBuilder builder = new StringBuilder();
		for (int col = 0; col < 4; col++) {
			for (int row = 0; row < 4; row++) {
				int index = col + 4 * row;
				byte b = get(index);
				String str = Integer.toHexString(0x000000FF & b);
				if (str.length() == 1)
					builder.append("0" + str + " ");
				else
					builder.append(str + " ");
			}
			builder.append("\n");
		}
		return builder.toString();
	}
	
	/*============================================================
	Private Methods
	============================================================*/
	
	/** Returns the byte at the given index in the Rijndael matrix. **/
	private final byte get(int index) {
		return state[offset + index];
	}

	/** Sets the byte at the given index of the Rijndael matrix to
	 * the given value. **/
	private final void set(int index, byte b) {
		state[offset + index] = b;
	}
	
	/** Performs a Galois field (256) multiplication of two bytes. **/
	public static byte gMul(byte x, byte y) {
		byte z = 0;
		for (byte i = 0; i < 8; i++) {
			if ((y & 1) != 0) {
				z ^= x;
			}
			final boolean hiBitSet = (x & 0x80) != 0;
			x <<= 1;
			if (hiBitSet) {
				x ^= 0x1b; /* x^8 + x^4 + x^3 + x + 1 */
			}
			y >>= 1;
		}
		return z;
	}
	
	/** Performs a Rcon operation on the given byte and returns the
	 * result. **/
	private static byte rcon(byte x) {
		final byte c2 = 0x02;
		byte c = 1;
		if (x == 0)
			return 0;
		while (x != 1) {
			c = gMul(c,c2);
			x--;
		}
		return c;
	}

}
